package chainx

import (
	"context"
	"gitee.com/dn-jinmin/langchaingo"
	"github.com/tmc/langchaingo/chains"
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/prompts"
	"strings"
)

const (
	_llmChainDefaultOutputKey = "text"
	LLMPromptString           = "prompt_string"
)

type LLMChain struct {
	*chains.LLMChain
	prompt prompts.FormatPrompter
}

var _ chains.Chain = &LLMChain{}

// NewLLMChain creates a new LLMChain with an LLM and a prompt.
func NewLLMChain(llm llms.Model, prompt prompts.FormatPrompter, opts ...chains.ChainCallOption) *LLMChain {
	return &LLMChain{
		LLMChain: chains.NewLLMChain(llm, prompt, opts...),
		prompt:   prompt,
	}
}

// Call formats the prompts with the input values, generates using the llm, and parses
// the output from the llm with the output parser. This function should not be called
// directly, use rather the Call or Run function if the prompt only requires one input
// value.
func (c LLMChain) Call(ctx context.Context, values map[string]any, options ...chains.ChainCallOption) (map[string]any, error) {
	out, promptStr, err := c.call(ctx, values, options...)
	if err != nil {
		return nil, err
	}

	return map[string]any{
		langchaingo.OutPutKey:    out,
		langchaingo.LLMPromptKey: promptStr,
		langchaingo.LLMOutPutKey: out,
	}, err
}

func (c LLMChain) call(ctx context.Context, values map[string]any, options ...chains.ChainCallOption) (string, string, error) {
	var promptStr string
	if v, ok := values[langchaingo.OutPutKey].(string); ok {
		promptStr = v
	} else {
		promptValue, err := c.Prompt.FormatPrompt(values)
		if err != nil {
			return "", promptStr, err
		}

		promptStr = promptValue.String()
	}

	result, err := llms.GenerateFromSinglePrompt(ctx, c.LLM, promptStr)
	if err != nil {
		return "", promptStr, err
	}

	return strings.TrimSpace(result), promptStr, nil
}
